import numpy as np
import torch
from sklearn.cluster import KMeans

from algorithms.algorithm import Algorithm
from algorithms.base.nsga2 import NSGA2
from config import device
from individual.population import Population
from operators.crossover.sbx import SBX
from operators.mutations.polynomial_mutation import PolynomialMutation
from problems.problem import Problem
from utils.vectors import get_ref_vectors, vector_radian, normalize_vector


class LMOEADS(Algorithm):
    """
    % Large-scale evolutionary multi-objective optimization assisted by
    % directed sampling
    %------------------------------- Reference --------------------------------
    % S. Qin, C. Sun, Y. Jin, Y. Tan, and J. Fieldsend, Large-scale
    % evolutionary multi-objective optimization assisted by directed sampling,
    % IEEE Transactions on Evolutionary Computation, 2021, 25(4): 724-738.
    %------------------------------- Copyright --------------------------------
    """
    parameters = [
        {'label': 'Nw', 'name': 'Nw', 'type': 'number', 'description': '种群聚类个数', 'step': 1, 'defaultValue': 10},
        {'label': 'Ns', 'name': 'Ns', 'type': 'number', 'step': 1, 'description': '沿每个引导方向的随机抽样次数',
         'defaultValue': 30},
    ]

    def __init__(self, problem: Problem,
                 pop_size: int,
                 kwargs,
                 ):
        """
        初始化算法
        :param problem:问题
        :param pop_size: 种群大小，实际大小会根据产生的权重向量个数决定
        :param max_func_eval: 最大评估次数
        :param kwargs: 其他参数
        """
        self.ref_vectors = get_ref_vectors(pop_size, problem.obj_dim)
        pop_size = len(self.ref_vectors)
        super().__init__(pop_size, problem, kwargs)
        self.cluster_vectors = torch.tensor(KMeans(n_clusters=self.Nw).fit(self.ref_vectors.cpu()).cluster_centers_,
                                            device=device)
        self.norm2_cluster_vectors = self.cluster_vectors.norm(dim=1, p=2)
        # 实际的种群大小
        pop_size = self.ref_vectors.shape[0]
        self.population = self.problem.init_population(pop_size)

    def each_iteration(self):
        obj = self.population.obj
        rep_solution_index = self.generate_representative_solution(self.population)
        # 论文中的SG，directed_sampling该方法返回了（offspring_dec,offspring_obj）元组
        guiding_population = self.directed_sampling(self.population, self.problem, rep_solution_index, self.Ns, self.Nw)
        self.population = self.double_reproduction(self.population, guiding_population, self.problem, self.ref_vectors,
                                                   self.pop_size)

    def generate_representative_solution(self, population):
        obj = population.obj
        radian_value = vector_radian(obj, self.cluster_vectors)
        obj_ref_min_index = radian_value.min(1)[1]
        unique_ref = torch.unique(obj_ref_min_index)
        # 保存代表解的个体索引
        rep_solution_index = []
        population_index = torch.arange(len(population), device=device)
        # 用于记录索引位置的individual是否已经再加入了集合population_temp中
        population_used = torch.zeros(obj.shape[0], dtype=bool)
        obj_min = obj.min(0)[0].unsqueeze(dim=0)
        # 在每个参考向量中选择一个最优的解
        for ref_vector in unique_ref:
            temp_index = obj_ref_min_index == ref_vector
            d1 = ((obj[temp_index] - obj_min) * self.cluster_vectors[ref_vector]).norm(p=2, dim=1) / \
                 self.norm2_cluster_vectors[ref_vector]
            min_d1_index = d1.argmin().item()
            rep_solution_index.append(population_index[temp_index][min_d1_index].item())
        # 设置该索引个体已经加入集合
        population_used[rep_solution_index] = True
        # 查看还有哪些参考向量没有分配个体
        no_individual_refs = np.setdiff1d(torch.arange(self.Nw).numpy(), unique_ref.numpy())
        if len(no_individual_refs) > 0:
            ref_obj_index = vector_radian(self.cluster_vectors[no_individual_refs], obj).sort(dim=1)[1]
            for i in range(len(no_individual_refs)):
                no_alloc_individual_index = ref_obj_index[i][population_used == False][0]
                rep_solution_index.append(population_index[no_alloc_individual_index].item())
                population_used[no_alloc_individual_index] = True
        return rep_solution_index

    @staticmethod
    def directed_sampling(population: Population, problem: Problem, rep_solution_index, Ns, Nw):
        """
        直接采样
        :param Nw:
        :param Ns:
        :param problem:
        :param population:
        :param rep_solution_index: 代表解的索引
        :return:
        """
        dec = population.dec[rep_solution_index]
        dec_low = dec - problem.low_limit
        dec_high = dec - problem.high_limit
        direction = torch.cat((dec_low, dec_high), dim=0)
        # 单位化direction
        direction = normalize_vector(direction, dim=1)
        # 计算决策变量上下界长度的平方和，再开根号
        interval = problem.interval_length.norm(p=2)
        # 直接采样
        rand_sample = torch.rand(Ns, 2 * Nw, device=device) * interval
        offspring_dec = []
        # 产生2*Ns*Nw个新的解
        for i in range(Ns):
            offspring_dec.append(
                problem.low_limit + rand_sample[i, :Nw].unsqueeze(dim=1) * direction[:Nw])
            offspring_dec.append(
                problem.high_limit + rand_sample[i, Nw:].unsqueeze(dim=1) * direction[Nw:])
        # 将list拼接成张量
        offspring_dec = torch.cat(offspring_dec, dim=0)
        offspring_dec = problem.repair_decision(offspring_dec)
        offspring = problem.estimate_population(offspring_dec)
        offspring_obj = offspring.obj
        # 非支配选择
        pf_solution_index = NSGA2.ND_sort(offspring_obj, 1)[0] == 1
        return Population(offspring_dec[pf_solution_index], offspring_obj[pf_solution_index])

    @classmethod
    def double_reproduction(cls, population: Population, guiding_population: Population, problem: Problem, ref_vectors,
                            pop_size):
        """
        两次production产生新的解
        :param ref_vectors:
        :param population:
        :param guiding_population:
        :param problem:
        :return:
        """
        population_dec = population.dec
        population_obj = population.obj
        guiding_solutions_dec = guiding_population.dec
        guiding_solutions_obj = guiding_population.obj
        # --------------------------First Phase start--------------------------------------
        # intermediate Q
        intermediate_Q = cls.first_reproduction(problem, population_dec, guiding_solutions_dec)
        # 三个种群dec合并（ P U intermediate_Q U S）--> intermediate_C
        intermediate_C = torch.cat((population_dec, intermediate_Q, guiding_solutions_dec), dim=0)
        # 三个种群obj合并
        intermediate_obj = torch.cat(
            (population_obj, problem.estimate_population(intermediate_Q).obj, guiding_solutions_obj), dim=0)
        # 大小和intermediate_obj的bool张量，用于选择intermediate_obj需要保留的解的索引
        next_index = cls.complementary_environmental_selection(intermediate_obj, ref_vectors, pop_size)
        # intermediate P
        intermediate_P, intermediate_obj = intermediate_C[next_index], intermediate_obj[next_index]
        # --------------------------First Phase end-------------------------------------

        # --------------------------Second Phase start-------------------------------------
        # 得到 offspring Ot
        # 把intermediate_P的大小凑成2的倍数
        if intermediate_P.size(0) % 2 == 1:
            intermediate_P = torch.cat((intermediate_P, intermediate_P[0].unsqueeze(dim=0)), dim=0)
            intermediate_obj = torch.cat((intermediate_obj, intermediate_obj[0].unsqueeze(dim=0)), dim=0)
        # 再次的交叉变异
        O = cls.second_reproduction(problem, intermediate_P[0::2], intermediate_P[1::2])
        O_offspring = problem.estimate_population(O)
        O_obj = O_offspring.obj
        # intermediate_P ∪ Ot
        obj = torch.cat((intermediate_obj, O_obj), dim=0)
        dec = torch.cat((intermediate_P, O), dim=0)
        nan_indices = torch.isnan(dec).any(dim=1)
        # 保留没有NaN的行
        dec = dec[~nan_indices]
        obj = obj[~nan_indices]
        next_index = cls.complementary_environmental_selection(obj, ref_vectors, pop_size)
        return Population(dec[next_index], obj[next_index])
        # --------------------------Second Phase end-------------------------------------

    @staticmethod
    def first_reproduction(problem, parent0, parent1, offspring_size=1):
        """
        将parent0, parent1交叉变异产生新的后代
        :param problem:
        :param parent0:
        :param parent1:
        :param offspring_size: 每一个将parent0，parent1交叉变异产生新的后代数量，只可以选1或者2
        :return: 返回一个[parent0.size(0)*offspring_size , parent0.size(1)]的新解集
        """
        # 随机排序种群
        rand_list = torch.randperm(parent0.size(0), device=device)

        # 采样得到的解的索引，随机从采样的解选择一个个体和种群的个体xxoo
        sampling_solution_index = torch.randint(parent1.size(0), (parent0.size(0),),
                                                device=device)
        parent0 = parent0[rand_list]
        parent1 = parent1[sampling_solution_index]

        offspring_dec = SBX.do(parent0, parent1, problem,
                               offspring_size=offspring_size)
        offspring_dec = PolynomialMutation.do(offspring_dec, problem=problem)
        return offspring_dec.unique(dim=0)

    @classmethod
    def complementary_environmental_selection(cls, intermediate_obj, ref_vectors, pop_size):
        """
        环境选择
        :param pop_size:
        :param ref_vectors:
        :param intermediate_obj:
        :return: 返回一个大小和intermediate_obj一样的bool矩阵
        """
        obj_max = intermediate_obj.max(0)[0]
        obj_min = intermediate_obj.min(0)[0]
        obj = (intermediate_obj - obj_min) / (obj_max - obj_min)
        obj_ref_max_value, obj_ref_max_index = vector_radian(obj, ref_vectors, return_radian=False).max(1)
        unique_index = obj_ref_max_index.unique()
        if unique_index.size(0) < (2 / 3) * pop_size:
            return cls.domination_selection(obj, pop_size)
        else:
            return cls.decomposition_selection(obj, obj_ref_max_index, obj_ref_max_value);

    @staticmethod
    def domination_selection(obj, pop_size):
        """
        支配选择
        :param pop_size:
        :param obj:
        :return: 返回一个大小和obj一样的bool矩阵
        """
        return NSGA2.environmental_selection(obj, pop_size)

    @staticmethod
    def decomposition_selection(obj, obj_ref_max_index, obj_ref_max_value):
        """
        分解选择
        :param obj:
        :param obj_ref_max_index: obj到最近参考向量的索引
        :param obj_ref_max_value: obj到最近参考向量的cos值
        :return:
        """
        obj = (obj - obj.min(0)[0]) / (obj.max(0)[0] - obj.min(0)[0])
        unique_index = torch.unique(obj_ref_max_index)
        next_index = torch.zeros(len(unique_index), dtype=torch.int)
        t = 0
        for index in unique_index:
            temp_index = torch.nonzero(index == obj_ref_max_index, as_tuple=True)[0]
            dist = torch.cdist(obj[temp_index], torch.zeros(1, obj.size(1), device=device, dtype=float)).squeeze(dim=1)
            fan = obj_ref_max_value[temp_index] / dist
            best = torch.argmax(fan)
            next_index[t] = temp_index[best]
            t += 1
        return next_index

    @classmethod
    def second_reproduction(cls, problem, patent0, patent1):
        return cls.first_reproduction(problem, patent0, patent1, offspring_size=2)
